from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch

# 1. 设置模型路径
model_path = r"D:/ideaSpace/MyPython/models/deepseek-coder-1.3b-base"

"""
使用量化配置会报错raise ImportError(
ImportError: Using `bitsandbytes` 8-bit quantization requires Accelerate: `pip install accelerate` and the latest version of bitsandbytes: `pip install -i https://pypi.org/simple/ bitsandbytes`
按要求安装accelerate和bitsandbytes也没有用
# 2. 量化配置 (8-bit 量化减少内存)
quant_config = BitsAndBytesConfig(
    load_in_4bit=True,          # 启用 4-bit 量化
    llm_int8_threshold=6.0      # 调节量化阈值
)
"""


# 3. 加载分词器（显式处理特殊标记）
tokenizer = AutoTokenizer.from_pretrained(
    model_path,
    pad_token="<pad>",
    eos_token="<|endoftext|>",
    add_special_tokens=False  # 关键参数
)

# 4. 加载量化模型（CPU）
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    device_map="cpu",
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True
)

# 5. 输入提示词
input_text = "def quick_sort(arr):"
inputs = tokenizer(input_text, return_tensors="pt").to("cpu")

# 6. 生成代码（限制长度以加速）
outputs = model.generate(
    **inputs,
    # do_sample=True,             # 启用采样，启用采用会使得生成结果更随机，temperature的值越大结果越随机（建议 0.5~1.0）
    max_length=100,            # 限制生成长度
    num_return_sequences=1     # 生成1个结果
    # temperature=0.7            # 控制随机性（temperature 仅当 do_sample=True（采样模式）时才会影响生成结果的随机性）
)

# 7. 解码并打印结果
print(tokenizer.decode(outputs[0], skip_special_tokens=True))